from comman.layers import *
from comman.losses import *
from comman.net import Sequential
from comman.optimizers import *
from dataset.captcha.captcha import load_captcha

epochs = 50
batch_size = 128
learning_rate = 1e-1

# 读入数据
(x_train, t_train), (x_test, t_test) = load_captcha(transpose=True)
x_validation, t_validation = x_test, t_test

net = Sequential([
    Conv2D(12, 3, (1, 1), activation='relu'),
    Pooling(pool_h=2, pool_w=2, stride=2),

    BatchNormalization(),
    Conv2D(36, (3, 5), (2, 3), activation='relu'),
    Pooling(pool_h=2, pool_w=2, stride=1),

    BatchNormalization(),
    Conv2D(128, (3, 5), (1, 1), activation='relu'),

    Flatten(),
    BatchNormalization(),
    Dense(128 * 2, activation=None),

    BatchNormalization(),
    Dense(4 * 36, activation="sigmoid"),
    Reshape([4, 36])
])

net.compile(MeanSquaredLoss(), SGD(learning_rate))

net.fit(x_train, t_train, epochs=epochs, batch_size=batch_size,
        validation_data=(x_validation, t_validation))

# accuracy = 0.9955
accuracy = net.evaluate(x_test, t_test)
print(f"accuracy={accuracy}")
